#include "net_rbf.h"
#include "node_bp.h"
#include "node_rbf.h"

#include "ann_math.h"

#include <time.h>

namespace ann
{

net_rbf::net_rbf(void): net(), m_hiCnt(0), m_w_yita(0.4), m_c_yita(0.5), m_s_yita(0.6)
{
    //hacking default struct
    m_lStruct.push_back(2);
    m_lStruct.push_back(2);
    m_lStruct.push_back(1);

    Init();
}

net_rbf::net_rbf(const UINT pInCnt, const UINT pHiCnt, const UINT pOutCnt): net(), m_hiCnt(0), m_w_yita(0.4), m_c_yita(0.5), m_s_yita(0.6)
{
    m_lStruct.push_back(pInCnt);
    m_lStruct.push_back(pHiCnt);
    m_lStruct.push_back(pOutCnt);
}

/*virtual*/ net_rbf::~net_rbf(void)
{
}

const UINT net_rbf::Init(const wchar_t* pInitFile /*= L"rbf_xor_init.info"*/)
{
    assert( m_lStruct.size() == 3 );

    m_inCnt  = m_lStruct[eRBF_In];
    m_hiCnt  = m_lStruct[eRBF_Hide];
    m_outCnt = m_lStruct[eRBF_Out];

    for( UINT i = 0; i < m_hiCnt; ++i )
    {
        node* rbfNode = new node_rbf(m_inCnt, &Gaussian);
        m_nodes.push_back(rbfNode);
    }
    
    for( UINT i = 0; i < m_outCnt; ++i )
    {
        node* bpNode = new node_bp(m_hiCnt, &Linear);
        m_nodes.push_back(bpNode);
    }

    // init default b, c and w
    InitFromFile(pInitFile);

    m_status |= eann_Initialized;
    return m_status;
}

const UINT net_rbf::Reset()
{
    m_status = 0;

    return m_status;   
}

const UINT net_rbf::Train()
{
    this->Elect_C_KNN();
    this->TRStaticCenter();

    return m_status;
}

const UINT net_rbf::Train_Gd()
{
    assert( m_outCnt == 1);

    const UINT sampCnt = m_trData.size();

    UINT counter = 0;

    double* dCt = new double[m_hiCnt*m_inCnt];
    double* dS  = new double[m_hiCnt];
    double* dW  = new double[m_hiCnt*m_outCnt];
    double* G   = new double[sampCnt*m_hiCnt];
    double* dis = new double[sampCnt*m_hiCnt];
    double* e   = new double[sampCnt];
    std::fill_n( dCt,  m_hiCnt*m_inCnt, 0.0);
    std::fill_n(  dS,          m_inCnt, 0.0);
    std::fill_n(  dW, m_hiCnt*m_outCnt, 0.0);
    std::fill_n(   G,  m_hiCnt*sampCnt, 0.0);
    std::fill_n( dis,  m_hiCnt*sampCnt, 0.0);
    std::fill_n(   e,          sampCnt, 0.0);

    double tVal = 0.0;
    double tSum_w = 0.0;
    double tSum_s = 0.0;
    double* tSum_c = new double[m_inCnt];
    std::fill_n(tSum_c, m_inCnt, 0.0);
    double wj = 0.0;
    double s = 0.0;

    for( ; !this->IsConverged(); )
    {
        ++counter;
        for( UINT i = 0; i < sampCnt; ++i )
        {
            TRData& curTRData = *m_trData[i];
            this->Sim(curTRData.m_inVec);
            curTRData.m_ek = this->CalEk(curTRData.m_outVec);
            //this->TRGD(curTRData.m_outVec, curTRData.m_inVec);
            //gain info per sample
            for( UINT j = 0; j < m_hiCnt; ++j )
            {
                  G[sampCnt*i+j] = m_nodes[GetNodeIdx(eRBF_Hide, j)]->Get_m_data();
                dis[sampCnt*i+j] = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, j)])->Get_m_dis();
            }
            e[i] = *curTRData.m_outVec.rbegin() - (*m_nodes.rbegin())->Get_m_data();
        }
        //update after all sample simulated
        // 1. first cal dw, dSigma, dCt
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            tSum_w = 0.0;
            for( UINT i = 0; i < sampCnt; ++i )
            {
                tSum_w += e[i] * G[sampCnt*i+j];
            }
            dW[j] = m_w_yita * tSum_w;   
        }
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            wj = dynamic_cast<node_bp*>(m_nodes[GetNodeIdx(eRBF_Out, 0)])->Get_m_weight(j);
            node_rbf* curNode = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, j)]);
            s  = curNode->Get_m_sigma();

            std::fill_n(tSum_c, m_inCnt, 0.0);
            tSum_s = 0.0;
            for( UINT i = 0; i < sampCnt; ++i )
            {
                tVal = e[i] * G[sampCnt*i+j];
                tSum_s += tVal * dis[sampCnt*i+j]*dis[sampCnt*i+j];
                for( UINT k = 0; k < m_inCnt; ++k )
                {
                    tSum_c[k] += tVal * (m_trData[i]->m_inVec[k] - curNode->Get_m_ctVec(k));
                }
            }
            dS[j] = m_s_yita * wj * tSum_s / (s*s*s);
            for( UINT k = 0; k < m_inCnt; ++k )
            {
                dCt[m_hiCnt*j+k] = m_c_yita * wj * tSum_c[k] / (s*s);
            }
        }
        // 2. apply dXX
        for( UINT j = 0; j < m_outCnt; ++j )
        {
            dW[j] = m_w_yita * tSum_w;   
            dynamic_cast<node_bp*>(m_nodes[GetNodeIdx(eRBF_Out, j)])->Get_m_weight(j) += dW[j];
        }
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            node_rbf* curNode = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, j)]);

            curNode->Get_m_sigma() += dS[j];
            for( UINT k = 0; k < m_inCnt; ++k )
            {
                curNode->Get_m_ctVec(k) += dCt[m_hiCnt*j+k];
            }
        }

        this->TestConverge();
    }
    
    if ( this->IsConverged() )
    {
        printf("%s: RBF net converged after %d iteration\n", __FUNCTION__, counter);
    }
    
    SAFE_DELA(dCt)
    SAFE_DELA(dS)
    SAFE_DELA(dW)
    SAFE_DELA(G)
    SAFE_DELA(dis)
    SAFE_DELA(e)
    SAFE_DELA(tSum_c)

    return m_status;
}

//this must be called after get train data
const UINT net_rbf::TRStaticCenter()
{
    assert( 1 == m_outCnt );

    const UINT GWidth = m_hiCnt+1;
    const UINT sampCnt= m_trData.size();

    //cal G and t
    double* G = new double[sampCnt*GWidth];
    double* t = new double[sampCnt];
    double* inVec = new double[m_inCnt];
    std::fill_n(G, sampCnt*GWidth, 0.0);
    std::fill_n(t, sampCnt, 0.0);
    std::fill_n(inVec, m_inCnt, 0.0);

    for( UINT i = 0; i < sampCnt; ++i )
    {
        TRData& curTRData = *m_trData[i];
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            node& curNode = *m_nodes[GetNodeIdx(eRBF_Hide,j)];
            std::copy(curTRData.m_inVec.begin(), curTRData.m_inVec.end(), inVec);
            G[GWidth*i+j] = curNode.tick(inVec, curTRData.m_inVec.size());
        }
        G[GWidth*i+GWidth-1] = 1.0;
        t[i] = *curTRData.m_outVec.begin();
    }

    //got G, then solve Gw = t
    double* w = new double[GWidth];
    std::fill_n(w, GWidth, 0.0);
    ann::math::SolveGwt(G, sampCnt, GWidth, t, w);

    std::copy(w, w+m_hiCnt, dynamic_cast<node_bp*>(*m_nodes.rbegin())->begin_m_weight());
    (*m_nodes.rbegin())->Set_m_b(w[GWidth-1]);

    node_bp* outNode = dynamic_cast<node_bp*>(*m_nodes.rbegin());
    std::cout<<"the weight & b is: ";
    std::for_each(outNode->begin_m_weight(), outNode->end_m_weight(), std::cout<<boost::lambda::_1<<", ");
    std::cout<<outNode->Get_m_b();
    std::cout<<std::endl;
    
    SAFE_DELA(G)
    SAFE_DELA(t)
    SAFE_DELA(inVec)
    SAFE_DELA(w)
    
    return m_status;
}

const UINT net_rbf::TRMSS()
{
    //var nodes, some mismatch with my struct
    //implement later

    return m_status;
}

//only works for mono output
const UINT net_rbf::TRGD(const dVec& pObjVec, const dVec& pInVec)
{
    assert( m_outCnt == 1 );
    assert( pObjVec.size() == m_outCnt );
    assert( pInVec.size()  == m_inCnt  );

    double* ek = new double[m_outCnt];
    std::fill(ek, ek+m_outCnt, 0.0);

    for( UINT i = 0; i < m_outCnt; ++i )
    {
        ek[i] = pObjVec[i] - m_nodes[GetNodeIdx(eRBF_Out, i)]->Get_m_data();
    }
    double e = ek[0];

    //adjust w
    for( UINT i = 0; i < m_outCnt; ++i )
    {
        node_bp* curNode = dynamic_cast<node_bp*>(m_nodes[GetNodeIdx(eRBF_Out, i)]);
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            curNode->Get_m_weight(j) += m_w_yita * e * m_nodes[GetNodeIdx(eRBF_Hide, j)]->Get_m_data();
        }
    }

    //adjust c and sigma
    double tVal = 0.0;
    for( UINT i = 0; i < m_hiCnt; ++i )
    {
        tVal = dynamic_cast<node_bp*>(m_nodes[GetNodeIdx(eRBF_Out, 0)])->Get_m_weight(i)*e;
        m_nodes[GetNodeIdx(eRBF_Hide, i)]->UpdateW(m_c_yita, tVal);
    }

    SAFE_DELA(ek)

    return m_status;
}

const UINT net_rbf::InitFromFile(const wchar_t* pInitFile)
{
    assert( pInitFile != NULL );

    std::ifstream inf( pInitFile, std::ifstream::in );
    if ( !inf.is_open() )
    {
        DEBUG("error: can't open init file.");
        m_status |= eann_FOINITFile;
        return m_status;
    }

    UINT inCnt = 0;
    UINT hiCnt = 0;
    UINT outCnt = 0;

    inf>>inCnt>>hiCnt>>outCnt;
    if ( inCnt!=m_inCnt || hiCnt!=m_hiCnt || outCnt!=m_outCnt )
    {
        DEBUG("error: init file error, net struct mismatch.");
        m_status |= eann_FEINITFile;
        return m_status;
    }

    double tVal = 0.0;

    for( UINT i = 0; i < m_hiCnt; ++i )
    {
        node*& curNode = m_nodes[GetNodeIdx(eRBF_Hide, i)];
        for( UINT j = 0; j < m_inCnt; ++j )
        {
            inf>>tVal;
            dynamic_cast<node_rbf*>(curNode)->Set_m_ctVec(tVal, j);
        }
        inf>>tVal;
        dynamic_cast<node_rbf*>(curNode)->Set_m_sigma(tVal);
        inf>>tVal;
        curNode->Set_m_b(tVal);
    }
    for( UINT i = 0; i < m_outCnt; ++i )
    {
        node*& curNode = m_nodes[GetNodeIdx(eRBF_Out, i)];
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            inf>>tVal;
            dynamic_cast<node_bp*>(curNode)->Set_m_weight(tVal, j);
        }
        inf>>tVal;
        curNode->Set_m_b(tVal);
    }

    return m_status;   
}

const UINT net_rbf::Sim(const dVec& pInVec, dVec* pOutVec /*= NULL*/)
{
    assert( pInVec.size() == m_inCnt );

    double* inVec = new double[m_inCnt];
    double* hiVec = new double[m_hiCnt];
    double* outVec= new double[m_outCnt];
    std::fill(inVec, inVec+m_inCnt, 0.0);
    std::fill(hiVec, hiVec+m_hiCnt, 0.0);
    std::fill(outVec, outVec+m_outCnt, 0.0);

    std::copy(pInVec.begin(), pInVec.end(), inVec);
    
    for( UINT i = 0; i < m_hiCnt; ++i )
    {
        node& curNode = *m_nodes[GetNodeIdx(eRBF_Hide,i)];
        hiVec[i] = curNode.tick(inVec, pInVec.size());
    }
    for( UINT i = 0; i < m_outCnt; ++i )
    {
        node& curNode = *m_nodes[GetNodeIdx(eRBF_Out,i)];
        outVec[i] = curNode.tick(hiVec, m_hiCnt);
    }

    if (pOutVec == NULL)
    {
        return m_status;
    }
    
    pOutVec->clear();
    for( UINT i = 0; i < m_outCnt; ++i )
    {
        pOutVec->push_back(outVec[i]);
    }

    SAFE_DELA(inVec)
    SAFE_DELA(hiVec)
    SAFE_DELA(outVec)

    return m_status;
}

const UINT net_rbf::TestConverge()
{
    double tVal = 0.0;

    for( trVec::iterator itr = m_trData.begin(); itr != m_trData.end(); ++itr )
    {
        tVal += (*itr)->m_ek;
    }
    tVal /= static_cast<double>(m_trData.size());
    tVal = sqrt(tVal);
    
    if ( tVal < m_criEk)
    {
        m_status |= eann_Converged;
    }

    DEBUG("%s: reached %2.8f\n", __FUNCTION__, tVal);
 
    return m_status;   
}

const UINT net_rbf::Elect_C_KNN()
{
    assert( m_trData.size() > m_hiCnt );

    using namespace boost;
    using namespace boost::lambda;

    // random init
    UINT tRand = 0;
    for( UINT i = 0; i < m_hiCnt; ++i )
    {
        node_rbf* hiNode = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, i)]);
        tRand = rand() % m_trData.size();
        std::copy(m_trData[tRand]->m_inVec.begin(), m_trData[tRand]->m_inVec.end(), hiNode->begin_m_ctVec());
    }

    double dmin = 10.0;
    double tVal = 0.0;
    double tSum = 1.0;
    UINT tCnt = 0;
    UINT* minDisIdx = new UINT[m_trData.size()];
    std::fill_n(minDisIdx, m_trData.size(), 0);
    double* nextC = new double[m_hiCnt*m_inCnt];
    std::fill_n(nextC, m_hiCnt*m_inCnt, 0.0);
    double* disAll = new double[m_hiCnt];
    std::fill_n(disAll, m_hiCnt, 0.0);
    UINT idx = 0;

#define EPSILON 0.0002
    for( ; tSum>EPSILON; )
    {
        std::fill_n(minDisIdx, m_trData.size(), 0);
        std::fill_n(nextC, m_hiCnt*m_inCnt, 0.0);

        // find min dis idx
        for( UINT i = 0; i < m_trData.size(); ++i )
        {
            dmin = 10.0;
            TRData& curTRData = *m_trData[i];
            this->Sim(curTRData.m_inVec);
            for( UINT j = 0; j < m_hiCnt; ++j )
            {
                disAll[j] = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, j)])->Get_m_dis();
            }
            minDisIdx[i] = std::min_element(disAll, disAll+m_hiCnt) - disAll;
            std::transform(curTRData.m_inVec.begin(), curTRData.m_inVec.end(), nextC+m_inCnt*minDisIdx[i], nextC+m_inCnt*minDisIdx[i], _1+_2);
        }

        //cal new center by average
        tSum = 0.0;
        for( UINT j = 0; j < m_hiCnt; ++j )
        {
            tVal = 0.0;
            node_rbf* node = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, j)]);
            tCnt = std::count(minDisIdx, minDisIdx+m_trData.size(), j);
            std::transform(nextC+m_inCnt*j, nextC+m_inCnt*(j+1), node->begin_m_ctVec(), node->begin_m_ctVec(), (ref(tVal) += (_1/tCnt-_2)*(_1/tCnt-_2), _1/tCnt));
            tSum += tVal;
        }
    }

    //cal sigma
    for( UINT j = 0; j < m_hiCnt; ++j )
    {
        tSum = 0.0;
        tCnt = std::count(minDisIdx, minDisIdx+m_trData.size(), j);
        node_rbf* node = dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide, j)]);
        for( UINT i = 0; i < m_trData.size(); ++i )
        {
            if ( minDisIdx[i] == j)
            {
                tVal = ann::math::EulerDis(m_trData[i]->m_inVec.begin(), m_trData[i]->m_inVec.end(), node->begin_m_ctVec());
                tSum += tVal*tVal;
            }
        }
        node->Set_m_sigma(sqrt(tSum/(2.0*tCnt)));

        printf("finally get the c for hide node %d is: ", j);
        std::for_each(node->begin_m_ctVec(), node->end_m_ctVec(), std::cout<<_1<<", ");
        std::cout<<std::endl;
    }

    return m_status;
}

const UINT net_rbf::InitStatic(const wchar_t* pInitStaticFile)
{
    assert( pInitStaticFile != NULL );

    std::ifstream inf( pInitStaticFile, std::ifstream::in );
    if ( !inf.is_open() )
    {
        DEBUG("error: can't open init file.");
        m_status |= eann_FOINITFile;
        return m_status;
    }

    UINT outCnt = 0;
    UINT hiCnt = 0;
    inf>>outCnt>>hiCnt;
    if ( outCnt!=m_outCnt || hiCnt!=m_hiCnt )
    {
        DEBUG("error: static init file error, net struct mismatch.");
        m_status |= eann_FEINITFile;
        return m_status;
    }

    for( UINT i = 0; i < m_outCnt; ++i )
    {
        node_bp* outNode = dynamic_cast<node_bp*>(m_nodes[GetNodeIdx(eRBF_Out, i)]);
        std::for_each(outNode->begin_m_weight(), outNode->end_m_weight(), inf>>boost::lambda::_1);
        inf>>(outNode->Get_m_b());
    }

    m_status |= eann_RBF_StaticInit;
    return m_status;
}

const UINT net_rbf::Elect_Sigma()
{
    const UINT sampCnt = m_trData.size();

    double dmax = 0.0;

    for( UINT i = 0; i < sampCnt-1; ++i )
    {
        for( UINT j = i+1; j < sampCnt; ++j )
        {
            double dis = ann::math::EulerDis(m_trData[i]->m_inVec.begin(), m_trData[i]->m_inVec.end(), m_trData[j]->m_inVec.begin());
            dmax = std::max(dmax, dis);
        }
    }

    const double sigma = dmax / sqrt(static_cast<double>(sampCnt));
    DEBUG("%s: find dmax = %f in %d samples, assign sigma = %f\n", __FUNCTION__, dmax, sampCnt, sigma);

    for( UINT i = 0; i < m_hiCnt; ++i )
    {
        dynamic_cast<node_rbf*>(m_nodes[GetNodeIdx(eRBF_Hide,i)])->Set_m_sigma(sigma);
    }

    return m_status;
}

}
